#! /usr/bin/env python3
'''
* add boolean in workflow inputs
* add if and scatter
'''


import argparse
import sys
import os
import re
import collections
bindir = os.path.abspath(os.path.dirname(__file__))
import logging
import WDL
import json
from pprint import pprint
import inspect


__author__='Liu Tao'
__mail__= 'taoliu@annoroad.com'

pat1=re.compile('^\s*$')

def log(infile):
    logger = logging.getLogger(__name__)
    logger.setLevel(level = logging.DEBUG)
    handler = logging.FileHandler(infile)
    handler.setLevel(logging.INFO)
    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    
    #logger.info("Start print log")
    #logger.debug("Do something")
    #logger.warning("Something maybe fail.")
    #logger.info("Finish")
    return(logger)

class Block():
    def __init__(self , obj ,region):
        self.obj = obj
        self.followed = []
        self.contains = []
        self.region = region
        self.shell = []
        self.inputs = {}
        self.outputs = {}
    def add_follow(self , i):
        self.followed.append(i)
    def add_contains(self ,i):
        self.contains.append(i)

class Workflow(Block):
    def __init__(self , obj ,region):
        super().__init__(obj ,region)
        self.name = obj.name
        self.id = "{0}|{1}".format(self.region , self.name )

    def add_wf_inputs( self , p_list , input_json , all_dict , traced_list) :  ###???
        for i in self.obj.inputs:  ### WDL.Tree.Decl
            name = i.name    
            value = get_decl_value( i , self.inputs ) #  , self.region , all_dict  )
            if len(p_list):
                p_obj = all_dict[ p_list[-1] ]
                if p_obj.inputs[name]:
                    value = p_obj.inputs[name]
                else:
                    print('using default value')
                    pass
            else:
                if name in input_json:
                    value = input_json[name]
                else:
                    pass
            self.inputs[name] = value
        pprint( [ 'inputs' , self.inputs , inspect.currentframe().f_lineno ])

    def add_wf_outputs(self, p_list , input_json , all_dict , traced_list):
        for o in self.obj.outputs:
            name = o.name
            if isinstance( o , WDL.Tree.Decl):
                if isinstance(o.expr , WDL.Expr.Get):
                    p_obj = all_dict[ p_list[-1] ]
                    value = get_get_key_and_value( o.expr  , a_dict=p_obj.inputs  ,  all_dict= all_dict , region= self.region)
                else:
                    print(o)
                    sys.exit("die in {0}".format(inspect.currentframe().f_lineno))
            self.outputs[name] = value
        pprint( ['output' , self.outputs , inspect.currentframe().f_lineno ])

def get_decl_value(obj , a_dict):
    if isinstance(obj , WDL.Tree.Decl):
        expr = obj.expr
        value = ''
        if isinstance( expr, WDL.Expr.String):
            value = get_string_key_and_value(expr , a_dict)
        elif isinstance(expr , WDL.Expr.Get):
            value = get_get_key_and_value(expr , a_dict)
        elif isinstance(expr , WDL.Expr.Apply):
            value = get_apply_key_and_value(expr , a_dict)
        elif isinstance(expr , WDL.Expr.Boolean):
            value = expr.value        #### may be use a function
        elif isinstance(expr , WDL.Expr.Array):
            value = get_array_key_and_value(expr , a_dict )
        elif isinstance(expr , WDL.Expr.Int):
            value = expr.value
        elif isinstance(expr , type(None)):
            value = ''
        else:
            print(type( expr ))
            raise Exception("die in {0}".format(inspect.currentframe().f_lineno))
    return(value)

def get_array_key_and_value(obj , a_dict):
    values = []
    for i in obj.items:
        if isinstance(i , WDL.Expr.Int):
            v = i.value
        else:
            print(type( i ))
            raise Exception("die in {0}".format(inspect.currentframe().f_lineno))
        values.append(v)
    return(values)

def get_apply_key_and_value(obj , a_dict):
    value = ''
    function_name = obj.function_name
    arguments = obj.arguments
    print('function is {0}'.format(function_name))
    values = []

    for i in arguments:
        if isinstance(i , WDL.Expr.String):
            value = get_string_key_and_value(i , a_dict)
        elif isinstance(i , WDL.Expr.Get):
            value = get_get_key_and_value(i , a_dict)
        values.append(value)

    if function_name == '_add':
        v = values[0] + values[1]
    return(v)

class Call(Block):
    def __init__(self , obj ,region ):
        super().__init__(obj ,region)
        self.name = obj.workflow_node_id
        self.id = "{0}|{1}".format(self.region , self.name )
        self.shell = []
        self.inputs['scatter_depth'] = {}

    def add_call_inputs(self , p_list , input_json , all_dict):
        for name in self.obj.inputs:
            a_io = self.obj.inputs[name]
            if isinstance(a_io , WDL.Expr.Get):
                #p_obj = all_dict[ p_list[-1] ]
                value = get_get_key_and_value( a_io  , a_dict=self.workflow.inputs  ,  all_dict= all_dict , region= self.region)
                if get_scatter_depth(a_io) > -1  : 
                    self.inputs['scatter_depth'][name] = get_scatter_depth(a_io)
            else:
                print(a_io)
                sys.exit("die in {0}".format(inspect.currentframe().f_lineno))
            #value = input_json[name]
            #print(value)
            self.inputs[name] = value
        pprint( [ 'inputs' , self.inputs , inspect.currentframe().f_lineno ])

    def add_call_outputs(self, p_list , input_json , all_dict):  #### continue in 2020-05-29
        for out in self.obj.callee.outputs:
            name = out.name
            value = ''
            if isinstance(out.expr , WDL.Expr.String ):
                value = get_string_key_and_value( out.expr , self.inputs )
                #print('=' * 40)
                #print(value)
            elif isinstance( out.expr , WDL.Expr.Get):
                region =  p_list[-1].split('|')[0]
                #p_obj = all_dict[ p_list[-1] ]
                value = get_get_key_and_value( out.expr  , a_dict=self.workflow.inputs  ,  all_dict= all_dict , region= region)
                #print('&' * 40)
                #print(value)
            else:
                print([ out , out.name , type(out.expr)])
                sys.exit("die in {0}".format(inspect.currentframe().f_lineno))
            ## directly from json
            self.outputs[name] = value   ### get placehold from inputs
        pprint( ['output' , self.outputs , inspect.currentframe().f_lineno ])

    def output_command(self , outdir ):
        if isinstance( self.obj.callee , WDL.Tree.Task):
            value = get_string_key_and_value( self.obj.callee.command , self.inputs)
            print(value)
        else:
            out = self.obj
            print([ out , out.name , type(out.expr)])
            sys.exit("die in {0}".format(inspect.currentframe().f_lineno))

        if isinstance( value , str):
            sh_file = '{0}/{1}_{2}.sh'.format(outdir , self.region , self.name )
            with open( sh_file , 'w') as f_out:
                f_out.write(value)
            #print(output)
            self.shell = [sh_file]
        else :
            for m,i in enumerate(value):
                sh_file = '{0}/{1}_{2}_{3}.sh'.format(outdir , self.region , self.name , m)
                with open( sh_file , 'w') as f_out:
                    f_out.write(i)
                #print(output)
                self.shell.append(sh_file)
        return(self.shell)

def get_string_key_and_value( obj , a_dict):
    v = []
    scatter_dict = {}
    for m , i in enumerate(obj.parts):
        if isinstance(i , str):
            v.append(i)
        elif isinstance(i , WDL.Expr.Placeholder):
            value = get_placeholder_key_and_value( i , a_dict)
            placehold_name = get_placeholder_name( i )
            if ('scatter_depth' in a_dict) and  (placehold_name in a_dict['scatter_depth']):
                scatter_dict[ a_dict['scatter_depth'][placehold_name] ] = m   #### return an error value ,  should use call.scatter_depth
            v.append(value)
            #print('^'*40)
        else:
            sys.exit("die in {0}".format(inspect.currentframe().f_lineno))
    if (len(scatter_dict)) > 0 : 
        r_values = []
        expand_all_line( r_values, v , scatter_dict )
        return(r_values)
    else:
        return("".join(v))

def get_placeholder_name( obj ):
    name = obj.expr.expr.referee.name
    return(name)

def expand_all_line( r_list , v , scatter_dict):
    key = sorted(scatter_dict.keys())[0]
    value = scatter_dict.pop(key)
    if len(scatter_dict) == 0:
        for i in v[value]:
            a_shell= "".join( v[ : value ] + [str(i)] + v[value+1 : ]  )
            r_list.append(a_shell)
    else:
        for i in v[value]:
            new_v = v[ : value] + [ str(i) ] + v[value+1 : ]
            expand_all_line(r_list , new_v , scatter_dict)

def get_placeholder_key_and_value( obj , a_dict):
    value = ''
    if isinstance(obj.expr , WDL.Expr.Get):
        value = get_get_key_and_value( obj.expr , a_dict =a_dict)
    else:
        sys.exit("die in {0}".format(inspect.currentframe().f_lineno))
    return(value)

def get_get_key_and_value(obj , a_dict = {} , all_dict={} , region='') : #,  region , all_dict , p_list , a_dict): 
    value = '' 
    key = ''
    expr = obj.expr 
    #print(type(expr.referee))
    if isinstance( expr.referee , WDL.Tree.Decl):
        key = expr.referee.name
        value = get_decl_value( expr.referee , a_dict ) 
        if key in a_dict :
            value = a_dict[key]
        else:
            pass
    elif isinstance(expr.referee , WDL.Tree.Call):    
        name = expr.referee.workflow_node_id   ### call name
        key = expr.name.split('.')[-1]
        if '{0}|{1}'.format(region , name) in all_dict: 
            if key in all_dict['{0}|{1}'.format(region , name)].outputs:
                value = all_dict['{0}|{1}'.format(region , name)].outputs[key]
                #print(value)
            else:
                pass
        else:
            pass
    elif isinstance(expr.referee , WDL.Tree.Scatter):
        key = expr.referee.variable
        name = expr.referee.workflow_node_id
        if '{0}|{1}'.format(region , name) in all_dict: 
            if key in all_dict['{0}|{1}'.format(region , name)].inputs:
                value = all_dict['{0}|{1}'.format(region , name)].inputs[key]    ### from there , scatter_depth
            else:
                pass
        else:
            pass
    else: 
        print(expr.referee)
        raise Exception("die in {0}".format(inspect.currentframe().f_lineno))
    return(value )

def get_scatter_depth(obj):
    expr = obj.expr
    if isinstance(expr.referee , WDL.Scatter):
        return (expr.referee.scatter_depth)
    else:
        return(-1)


class Task(Block):
    def __init__(self, obj ,region):
        super().__init__(obj ,region)
        self.name = obj.name
        self.cmd = obj.command.parts

class If(Block):
    def __init__(self , obj ,region):
        super().__init__(obj ,region)
        self.name = obj.workflow_node_id
        self.variable = []
        self.function = []
        self.id = "{0}|{1}".format(self.region , self.name )
    def add_if_inputs( self, p_list , input_json , all_dict ):
        expr = self.obj.expr
        if isinstance(expr , WDL.Expr.Get):
            self.value = get_get_key_and_value( expr  , a_dict=self.workflow.inputs  ,  all_dict= all_dict , region= self.region)
        else:
            raise Exception("die in {0}".format(inspect.currentframe().f_lineno))

class Scatter(Block):
    def __init__(self , obj ,region):
        super().__init__(obj ,region)
        self.name = obj.workflow_node_id
        self.variable = []
        self.function = []
        self.id = "{0}|{1}".format(self.region , self.name )
    def add_scatter_inputs( self ,  previous_list , input_json , all_dict ):
        name = self.obj.variable
        expr = self.obj.expr
        if isinstance(expr , WDL.Expr.Get):
            value = get_get_key_and_value(expr , a_dict=self.workflow.inputs  ,  all_dict= all_dict , region= self.region)
        else:
            print(a_io)
            sys.exit("die in {0}".format(inspect.currentframe().f_lineno))
        self.inputs[name] = value
        pprint( [ 'inputs' , self.inputs , inspect.currentframe().f_lineno ])

def trace(doc , a_dict , region):  
    if isinstance(doc , WDL.Document):
        if doc.workflow:
            wf = Workflow(  doc.workflow , region )
            a_dict[ wf.id ] = wf
            trace( wf , a_dict , region) 
    elif isinstance(doc.obj , WDL.Workflow):
        for i in doc.obj.body:
            if isinstance( i , WDL.Decl ):
                pass
            elif isinstance( i , WDL.Call):
                call = Call( i , region )
                call.workflow = doc  #### use for get_node later on 
                a_dict[ call.id ] = call
                doc.add_contains( call.id )
                trace( call ,  a_dict ,region)
            elif isinstance(i , WDL.Scatter ):
                scatter = Scatter( i , region )
                scatter.workflow = doc
                a_dict[scatter.id] = scatter
                doc.add_contains( scatter.id )
                trace( scatter , a_dict , region)
            elif isinstance(i , WDL.Conditional ):
                condition = If( i , region )
                condition.workflow = doc
                a_dict[condition.id] = condition
                doc.add_contains( condition.id )
                trace( condition , a_dict , region)
            else:
                pprint(i.__dict__)
                pprint(i)
                sys.exit("die in {0}".format(inspect.currentframe().f_lineno))
    elif isinstance(doc.obj , WDL.Call):
        if isinstance( doc.obj.callee , WDL.Workflow):
            region = region + '-' + doc.obj.callee_id[0]
            wf = Workflow(  doc.obj.callee , region )
            a_dict[ wf.id ] = wf
            doc.add_contains( wf.id )
            trace(wf , a_dict ,region )
    elif isinstance(doc.obj , WDL.Task):
        pass
    elif isinstance(doc.obj , WDL.Scatter):
        for i in doc.obj.body:
            if isinstance( i , WDL.Decl ):
                pass
            elif isinstance( i , WDL.Call):
                call = Call( i , region )
                call.workflow = doc.workflow   #### use for get_node later on 
                a_dict[ call.id ] = call
                doc.add_contains( call.id )
                trace( call ,  a_dict ,region)
            elif isinstance(i , WDL.Scatter ):
                scatter = If( i , region )
                scatter.workflow = doc.workflow 
                a_dict[scatter.id] = scatter
                doc.add_contains( scatter.id )
                trace( scatter , a_dict , region)
            elif isinstance(i , WDL.Conditional ):
                condition = If( i , region )
                condition.workflow = doc.workflow 
                a_dict[condition.id] = condition
                doc.add_contains( condition.id )
                trace( condition , a_dict , region)
            else:
                pprint(i.__dict__)
                pprint(i)
                raise Exception("die in {0}".format(inspect.currentframe().f_lineno))
    elif isinstance(doc.obj , WDL.Conditional):
        for i in doc.obj.body:
            if isinstance( i , WDL.Decl ):
                pass
            elif isinstance( i , WDL.Call):
                call = Call( i , region )
                call.workflow = doc.workflow   #### use for get_node later on 
                a_dict[ call.id ] = call
                doc.add_contains( call.id )
                trace( call ,  a_dict ,region)
            elif isinstance(i , WDL.Scatter ):
                pass
            elif isinstance(i , WDL.Conditional ):
                condition = If( i , region )
                condition.workflow = doc.workflow 
                a_dict[condition.id] = condition
                doc.add_contains( condition.id )
                trace( condition , a_dict , region)
            else:
                pprint(i.__dict__)
                pprint(i)
                raise Exception("die in {0}".format(inspect.currentframe().f_lineno))
    else:
        pprint(doc)
        pprint(doc.__dict__)
        sys.exit("die in {0}".format(inspect.currentframe().f_lineno))

def get_all_path(start , a_dict , path , total_path ):
    obj = a_dict[start]
    path.append(start)
    if len(obj.contains) > 1 :
        for i in obj.contains:
            new_path = path[:]
            get_all_path(i , a_dict , new_path , total_path)
    elif len(obj.contains) == 1:
        i = obj.contains[0]
        get_all_path(i , a_dict , path ,total_path)
    else:
        total_path.append(path)

def sort_all_path(all_path , obj_dict):
    '''
    add dependency obj into the all_path , and sort them 
    '''
    new_path_list = []
    found_set = set()
    while(all_path):
        a_path = all_path.pop(0)
        depend_list = get_dependency(a_path , obj_dict)
        #print(a_path)
        #print(depend_list)
        #print(found_set)  ### incluse workflow name 
        if len(depend_list):
            if set(depend_list).issubset(found_set):
                new_path_list.append([a_path ,  depend_list])
                found_set.update(set(a_path))
            else:
                all_path.append(a_path)
        else:
            new_path_list.append([a_path ,  depend_list])
            found_set.update(set(a_path))
    return(new_path_list)

def get_dependency(a_path ,a_dict):
    '''
    get dependency  
    '''
    r_list = []
    for i in a_path:
        obj = a_dict[i]
        if isinstance(obj , Workflow):
            pass
        elif isinstance(obj , Call) or isinstance(obj , If) or isinstance(obj , Scatter):
            #pprint([i, obj.__dict__])
            trace_dependecies(obj.obj.workflow_node_dependencies , obj.workflow.obj , obj.region,  r_list)
        else:
            raise Exception(i , obj)
    return(r_list)

def trace_dependecies(dep_set , wf , region , r_list):
    for i in dep_set:
        new_node = wf.get_node(i)
        if isinstance(new_node , WDL.Call):
            name = '{0}|{1}'.format(region , new_node.workflow_node_id)
            r_list.append( name )   ##### this name should be careful inspect
        new_dep = new_node.workflow_node_dependencies
        if len(new_dep):
            trace_dependecies(new_dep , wf , region ,  r_list)
        else:
            pass

def output_shell(all_sorted_path , obj_dict, input_json , outdir ):
    traced_list = []
    dep_sh = '{0}/dependent.list'.format(outdir)
    f_out = open(dep_sh,'w')
    for i in all_sorted_path:
        print('+'*50)
        previous_list = [] 
        bool_if = True
        for a_obj_name in i[0]:
            obj = obj_dict[ a_obj_name ]
            print(['****' , bool_if ,  obj.id , obj, previous_list])
            if isinstance(obj , Workflow):
                obj.add_wf_inputs(  previous_list , input_json , obj_dict , traced_list)
            elif isinstance(obj , Call):
                obj.add_call_inputs( previous_list , input_json , obj_dict  )
            elif isinstance(obj , If):
                obj.add_if_inputs( previous_list , input_json , obj_dict )
                if bool_if : bool_if = obj.value
            elif isinstance(obj , Scatter):
                obj.add_scatter_inputs( previous_list , input_json , obj_dict )
                pass
            else:
                raise Exception( "die in {0}".format(inspect.currentframe().f_lineno))
            previous_list.append( a_obj_name )
            traced_list.append( a_obj_name )

        previous_list = [] 
        for a_obj_name in reversed(i[0]):
            obj = obj_dict[ a_obj_name ]
            print(['****' , obj.id , obj, previous_list])
            if isinstance(obj , Workflow):
                obj.add_wf_outputs(  previous_list , input_json , obj_dict , traced_list)
            elif isinstance(obj , Call):
                obj.add_call_outputs( previous_list , input_json , obj_dict  )
            elif isinstance(obj , If):
                pass
            elif isinstance(obj, Scatter):
                ## may be add an output
                pass
            else:
                raise Exception( "die in {0}".format(inspect.currentframe().f_lineno))
            previous_list.append( a_obj_name )
        
        a_shell = []
        for count , a_obj_name in enumerate(reversed(i[0])):
            obj = obj_dict[ a_obj_name ]
            if count == 0 : 
                a_shell = obj.output_command( outdir )
            else:
                obj.shell += a_shell
        if bool_if:
            for a_obj_name in i[1]:
                for a_dep_sh in obj_dict[a_obj_name].shell:
                    for a_sh in a_shell:
                        f_out.write('{0}\t{1}\n'.format(a_dep_sh , a_sh))
            if len( i[1] ) == 0 :
                for a_sh in a_shell:
                    f_out.write('{0}\n'.format(a_sh))
    f_out.close()

def main():
    parser=argparse.ArgumentParser(description=__doc__,
            formatter_class=argparse.RawDescriptionHelpFormatter,
            epilog='author:\t{0}\nmail:\t{1}'.format(__author__,__mail__))
    parser.add_argument('-i','--input',help='input file',dest='input',required=True)
    parser.add_argument('-j','--json',help='input json file',dest='json', type=open, required=True)
    parser.add_argument('-o','--outdir',help='output shell directory',dest='outdir',required=True)
    args=parser.parse_args()

    global logger 
    logfile = '{0}/log.txt'.format(args.outdir)
    logger = log( logfile )

    doc = WDL.load(args.input)

    obj_dict = collections.OrderedDict()
    logger.info("Start print log")
    trace(doc , obj_dict , 'main')
    pprint(obj_dict)

    first_obj_name = list(obj_dict.keys())[0]

    all_path = []
    get_all_path(first_obj_name , obj_dict , [] , all_path )
    pprint(all_path)

    all_sorted_path = sort_all_path(all_path , obj_dict)
    pprint(all_sorted_path)

    input_json = json.load(args.json)

    output_shell(all_sorted_path , obj_dict , input_json , args.outdir )

    print('done')


if __name__ == '__main__':
    main()

